import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
from scipy.integrate import odeint

# 定义3×3×3网格参数
size = 3
num_nodes = size**3
k = 1.0       # 弹簧刚度
m = 1.0       # 节点质量
damping = 0.1 # 阻尼系数

# 初始化节点位置（三维坐标）
nodes = np.array([(x, y, z) for x in range(size) for y in range(size) for z in range(size)])

# 构建邻居连接（每个节点连接上下左右前后6个方向）
neighbors = {}
for i in range(num_nodes):
    x, y, z = nodes[i]
    neighbors[i] = []
    # 检查六个方向的相邻节点是否存在
    for dx, dy, dz in [(-1,0,0), (1,0,0), (0,-1,0), (0,1,0), (0,0,-1), (0,0,1)]:
        nx, ny, nz = x + dx, y + dy, z + dz
        if 0 <= nx < size and 0 <= ny < size and 0 <= nz < size:
            j = nx * size**2 + ny * size + nz
            neighbors[i].append(j)

# 定义运动方程
def equations(y, t):
    displacements = y[:3*num_nodes].reshape(num_nodes, 3)
    velocities = y[3*num_nodes:].reshape(num_nodes, 3)
    dydt = np.zeros_like(y)
    
    # 计算每个节点的加速度
    for i in range(num_nodes):
        force = np.zeros(3)
        # 邻居节点施加的弹簧力
        for j in neighbors[i]:
            delta = displacements[j] - displacements[i]
            distance = np.linalg.norm(nodes[j] - nodes[i])
            force += k * delta / distance  # 胡克定律
        # 阻尼力
        force -= damping * velocities[i]
        # 更新加速度
        dydt[3*num_nodes + i*3 : 3*num_nodes + (i+1)*3] = force / m
    # 更新速度
    dydt[:3*num_nodes] = velocities.flatten()
    return dydt

# 初始条件：中间节点施加初始位移
initial_displacements = np.zeros((num_nodes, 3))
initial_velocities = np.zeros((num_nodes, 3))
center_idx = size**3 // 2  # 中心节点索引
initial_displacements[center_idx] = [0.5, 0, 0]  # x方向位移

# 合并初始条件并求解微分方程
t = np.linspace(0, 10, 100)
y0 = np.concatenate([initial_displacements.flatten(), initial_velocities.flatten()])
solution = odeint(equations, y0, t)

# 提取最后一个时间步的位移结果
final_displacements = solution[-1, :3*num_nodes].reshape(num_nodes, 3)

# 3D可视化
fig = plt.figure(figsize=(10, 8))
ax = fig.add_subplot(111, projection='3d')
ax.scatter(nodes[:,0], nodes[:,1], nodes[:,2], c='blue', label='oringinal position') #原始位置
ax.scatter(nodes[:,0] + final_displacements[:,0], 
           nodes[:,1] + final_displacements[:,1], 
           nodes[:,2] + final_displacements[:,2], 
           c='red', label='after vibration') #振动后位置
ax.set_xlabel('X')
ax.set_ylabel('Y')
ax.set_zlabel('Z')
# plt.scatter(x, y, c=categories, cmap='viridis')
# plt.legend(['Blue Category', 'Red Category'])  # 明确的图例标签
plt.legend()
plt.show()